"""
训练和验证逻辑模块
包含模型训练和验证的核心逻辑
"""
from models.dataset import ImageFromURLDataset, get_model_transform
from models.model_utils import get_device, get_optimal_workers, create_resnext_model, save_model_and_labels, load_model_and_labels, cleanup_gpu_memory
from utils.file_utils import cleanup_temp_files

def create_dataset_with_torch_inheritance(json_path, transform=None):
    """创建具有torch Dataset继承的数据集"""
    # 只继承ImageFromURLDataset，避免MRO冲突
    class TorchImageFromURLDataset(ImageFromURLDataset):
        def __init__(self, json_path, transform=None):
            super().__init__(json_path, transform)
    return TorchImageFromURLDataset(json_path, transform)


def run_training(json_path, model_folder, socketio, current_model_name_callback):
    """在一个单独的线程中运行模型训练"""
    # 延迟导入torch相关模块
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader
    
    socketio.emit('status_update', {'msg': '训练准备中：正在初始化数据集...'})
    
    # 1. 数据准备
    transform = get_model_transform(is_training=True)

    try:
        dataset = create_dataset_with_torch_inheritance(json_path=json_path, transform=transform)
        num_classes = dataset.get_num_classes()
        if num_classes == 0:
            socketio.emit('status_update', {'msg': '错误：JSON文件中未找到有效的图片或标签。'})
            return

        # 在 macOS 上使用 num_workers=0 避免多进程资源泄漏
        num_workers = get_optimal_workers()
        dataloader = DataLoader(dataset, batch_size=16, shuffle=True, num_workers=num_workers)
        
    except Exception as e:
        socketio.emit('status_update', {'msg': f'错误：无法加载数据集。 {e}'})
        return

    # 2. 模型设置
    socketio.emit('status_update', {'msg': f'训练准备中：正在加载ResNeXt模型，共 {num_classes} 个分类...'})
    
    device = get_device()
    model = create_resnext_model(num_classes, device)

    if device.type == 'cuda' and torch.cuda.device_count() > 1:
        print(f"检测到 {torch.cuda.device_count()} 张GPU，启用DataParallel多卡训练。")
        model = torch.nn.DataParallel(model)
        model = model.to(device)

    if device.type == 'mps':
        print("MPS is available:", torch.backends.mps.is_available())
        print("MPS is built:", torch.backends.mps.is_built())
        print("Torch version:", torch.__version__)

    # 3. 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 4. 训练循环
    num_epochs = 5  #  epochs

    total_steps = len(dataloader)



    print("##START ： run_training###") # 按照用户要求，在训练循环前添加调试输出
    epoch = 0

    socketio.emit('status_update', {'msg': f'训练开始！设备: {device}, 轮次: {num_epochs}/{epoch}'})

    for epoch in range(num_epochs):

        # 清理25分钟前的临时jpg文件
        cleanup_temp_files('/data/temp/images', 25)

        model.train()
        running_loss = 0.0
        running_corrects = 0
        total_samples = 0

        for i, (inputs, labels) in enumerate(dataloader):
            try:
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                with torch.set_grad_enabled(True):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    loss.backward()
                    optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                total_samples += inputs.size(0)
                
                # 发送实时数据到前端
                current_loss = running_loss / total_samples
                current_acc = running_corrects.double() / total_samples * 100
                
                socketio.emit('update_stats', {
                    'epoch': epoch + 1,
                    'step': i + 1,
                    'total_steps': total_steps,
                    'loss': current_loss,
                    'accuracy': current_acc.item()
                })
                socketio.sleep(0.01) # 给予服务器处理其他请求的时间
                
            except Exception as e:
                socketio.emit('status_update', {'msg': f'训练步骤出错: {e}'})
                continue

        # 清理GPU/MPS内存
        cleanup_gpu_memory(device)
            
        epoch_loss = running_loss / len(dataloader.dataset)
        epoch_acc = running_corrects.double() / len(dataloader.dataset) * 100
        
        socketio.emit('status_update', {
            'msg': f' 设备：{device} | 轮次 {epoch + 1}/{num_epochs} 完成 | 损失: {epoch_loss:.4f} | 准确率: {epoch_acc:.2f}%'
        })

    # --- 保存带时间戳的模型和标签 ---
    try:
        model_filename = save_model_and_labels(model, dataset.get_labels(), model_folder)
        
        # 更新当前模型名称
        current_model_name_callback(model_filename)
        socketio.emit('training_complete', {
            'msg': f'训练已成功完成！模型已保存为 {model_filename}',
            'model_name': model_filename
        })
    except Exception as e:
        socketio.emit('status_update', {'msg': f'保存模型时出错: {e}'})


def run_validation(json_path, model_folder, current_model_name, socketio):
    """在一个单独的线程中运行模型验证"""
    # 延迟导入torch相关模块
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader
    
    print("##START ： run_training###") # 按照用户要求，在训练循环前添加调试输出

    if not current_model_name:
        socketio.emit('status_update', {'msg': '错误：当前没有可用的模型。请先训练或上传一个模型。'})
        return

    socketio.emit('status_update', {'msg': '验证准备中：正在加载模型和数据...'})
    
    device = get_device()



    try:
        model, labels_list = load_model_and_labels(current_model_name, model_folder, device)
        model.eval()
        num_classes = len(labels_list)
   
    except Exception as e:
        socketio.emit('status_update', {'msg': f'错误：无法加载模型。 {e}'})
        return

    # 1. 数据准备
    transform = get_model_transform(is_training=True)

    try:
        dataset = create_dataset_with_torch_inheritance(json_path=json_path, transform=transform)
        if num_classes != dataset.get_num_classes():
             socketio.emit('status_update', {'msg': '错误：验证数据的类别数量与已训练模型的类别数量不匹配。'})
             return
        num_workers = get_optimal_workers()
        dataloader = DataLoader(dataset, batch_size=16, shuffle=False, num_workers=num_workers)
    except Exception as e:
        socketio.emit('status_update', {'msg': f'错误：无法加载验证数据集。 {e}'})
        return

    criterion = nn.CrossEntropyLoss()
    
    socketio.emit('status_update', {'msg': f'验证开始！设备: {device}'})
    # 发送当前使用的模型名称
    socketio.emit('current_model_info', {'model_name': current_model_name})
    
    running_loss = 0.0
    running_corrects = 0
    total_samples = 0
    total_steps = len(dataloader)

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloader):
            try:
                inputs = inputs.to(device)
                labels = labels.to(device)

                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                total_samples += inputs.size(0)

                current_loss = running_loss / total_samples
                current_acc = running_corrects.double() / total_samples * 100
                
                socketio.emit('update_validation_stats', {
                    'step': i + 1,
                    'total_steps': total_steps,
                    'loss': current_loss,
                    'accuracy': current_acc.item()
                })
                socketio.sleep(0.01)
                
            except Exception as e:
                socketio.emit('status_update', {'msg': f'验证步骤出错: {e}'})
                continue

    # 清理GPU/MPS内存
    cleanup_gpu_memory(device)

    final_loss = running_loss / len(dataset)
    final_acc = running_corrects.double() / len(dataset) * 100
    
    socketio.emit('validation_complete', {
        'msg': f'验证完成 | 总损失: {final_loss:.4f} | 总准确率: {final_acc:.2f}%'
    })
